from datetime import datetime
from py2neo import Graph, NodeMatcher, RelationshipMatcher

class Neo4jMap:
    def __init__(self):
        # Neo4j连接配置
        self.neo_url = "bolt://localhost:7687"
        self.neo_user = "neo4j"
        self.neo_password = "324820"
        self.ggraph = ""
        self.node_matcher = ""
        self.relationship_matcher = ""

    def neo4j_connect(self):
        # 初始化Neo4j连接
        try:
            # 实例化图数据库
            # 连接neo4j数据库，输入地址、用户名、密码
            self.graph = Graph(self.neo_url, auth=(self.neo_user, self.neo_password))
            # 实例化节点查询对象
            self.node_matcher = NodeMatcher(self.graph)
            # 实例化关系查询对象
            self.relationship_matcher = RelationshipMatcher(self.graph)
        except Exception as e:
            print(f"数据库连接失败: {e}")

    def node_to_dict(self, node):
        """将Py2neo节点转换为字典"""
        if not node:
            return None
        return {
            "identity": node.identity,
            "labels": list(node.labels),
            "properties": dict(node)
        }

    def relationship_to_dict(self, rel):
        """将Py2neo关系转换为字典"""
        return {
            "identity": rel.identity,
            "start": rel.start_node.identity,
            "end": rel.end_node.identity,
            "type": rel.__class__.__name__,
            "properties": dict(rel)
        }

    def get_all_spo(self):
        """获取数据库中所有的SPO关系"""
        try:
            # 查询所有的关系
            query = """
            MATCH (start)-[rel]->(end)
            RETURN start, rel, end
            ORDER BY start.name, rel.__type__, end.name
            LIMIT 1000
            """
            # 执行查询并获取数据
            results = self.graph.run(query).data()
            # 定义关系数据列表
            spo_list = []
            # 处理后的关系字典
            processed_relationships = set()
            # 遍历查询结果
            for record in results:
                start_node = record['start']
                rel = record['rel']
                end_node = record['end']
                # 避免重复关系
                rel_key = f"{start_node.identity}-{rel.identity}-{end_node.identity}"
                if rel_key in processed_relationships:
                    continue
                processed_relationships.add(rel_key)
                spo_data = {
                    "start": self.node_to_dict(start_node),
                    "end": self.node_to_dict(end_node),
                    "relationship": self.relationship_to_dict(rel),
                    "segments": [
                        {
                            "start": self.node_to_dict(start_node),
                            "end": self.node_to_dict(end_node),
                            "relationship": self.relationship_to_dict(rel),
                        }
                    ],
                    "length": 1.0
                }
                spo_list.append(spo_data)
            count = len(spo_list)
            timestamp = datetime.now().isoformat()
            return spo_list, count, timestamp
        except Exception as e:
            return e

    def get_filtered_spo(self, start_label, rel_type, end_label, search_term=""):
        """根据条件过滤SPO关系"""
        try:
            # 构建查询
            query_parts = ["MATCH (start)-[rel]->(end)"]
            where_conditions = []
            params = {}
            # 处理节点标签过滤 - 使用正确的语法
            if start_label:
                where_conditions.append("ANY(label IN labels(start) WHERE label = $start_label)")
                params['start_label'] = start_label
            if end_label:
                where_conditions.append("ANY(label IN labels(end) WHERE label = $end_label)")
                params['end_label'] = end_label
            # 处理关系类型过滤
            if rel_type:
                where_conditions.append("type(rel) = $rel_type")
                params['rel_type'] = rel_type
            # 处理搜索词过滤
            if search_term:
                where_conditions.append("(start.name CONTAINS $search_term OR end.name CONTAINS $search_term)")
                params['search_term'] = search_term
            if where_conditions:
                query_parts.append("WHERE " + " AND ".join(where_conditions))
            query_parts.append("RETURN start, rel, end")
            query_parts.append("ORDER BY start.name, type(rel), end.name")
            query_parts.append("LIMIT 500")
            query = " ".join(query_parts)
            results = self.graph.query(query, params)
            spo_list = []
            for record in results:
                start_node = record['start']
                rel = record['rel']
                end_node = record['end']
                spo_data = {
                    "start": self.node_to_dict(start_node),
                    "end": self.node_to_dict(end_node),
                    "relationship": self.relationship_to_dict(rel),
                    "segments": [
                        {
                            "start": self.node_to_dict(start_node),
                            "end": self.node_to_dict(end_node),
                            "relationship": self.relationship_to_dict(rel),
                        }
                    ],
                    "length": 1.0
                }
                spo_list.append(spo_data)
            count = len(spo_list)
            print(spo_list)
            timestamp = datetime.now().isoformat()
            return count, spo_list, query, timestamp
        except Exception as e:
            return e

    def get_metadata(self):
        """获取数据库元数据"""
        try:
            # 获取所有节点标签
            labels_query = "MATCH (n) RETURN DISTINCT labels(n) AS labels"
            labels_result = self.graph.run(labels_query).data()
            all_labels = set()
            for record in labels_result:
                all_labels.update(record['labels'])
            # 获取所有关系类型
            rel_types_query = "MATCH ()-[r]->() RETURN DISTINCT type(r) AS rel_type"
            rel_types_result = self.graph.run(rel_types_query).data()
            all_rel_types = {record['rel_type'] for record in rel_types_result}

            # 获取统计信息
            stats_query = """
            MATCH (n)
            WITH count(n) AS node_count
            MATCH ()-[r]->()
            RETURN node_count, count(r) AS relationship_count
            """
            stats_result = list(self.graph.run(stats_query).data())
            labels = sorted(list(all_labels))
            relationship_types = sorted(list(all_rel_types))
            node_count = stats_result[0]['node_count'] if len(stats_result) > 0 else 0
            relationship_count = stats_result[1]['relationship_count'] if len(stats_result) > 1 else 0
            return labels, relationship_types, node_count, relationship_count
        except Exception as e:
            return e

    def get_visualization_data(self, limit):
        """获取可视化数据"""
        try:
            query = f"""
            MATCH (start)-[rel]->(end)
            RETURN start, rel, end
            LIMIT {limit}
            """
            results = self.graph.run(query).data()
            nodes = []
            edges = []
            node_ids = set()
            for record in results:
                start_node = record['start']
                rel = record['rel']
                end_node = record['end']
                # 处理起始节点
                if start_node.identity not in node_ids:
                    nodes.append({
                        'id': start_node.identity,
                        'label': start_node.get('name', list(start_node.labels)[0])
                                    if start_node.labels else 'Unknown',
                        'group': list(start_node.labels)[0] if start_node.labels else 'Unknown',
                        'title': f"{list(start_node.labels)[0] if start_node.labels else 'Unknown'}:"
                                 f" {start_node.get('name', 'No name')}",
                        'properties': dict(start_node)
                    })
                    node_ids.add(start_node.identity)

                # 处理目标节点
                if end_node.identity not in node_ids:
                    nodes.append({
                        'id': end_node.identity,
                        'label': end_node.get('name', list(end_node.labels)[0]) if end_node.labels else 'Unknown',
                        'group': list(end_node.labels)[0] if end_node.labels else 'Unknown',
                        'title': f"{list(end_node.labels)[0] if end_node.labels else 'Unknown'}: "
                                 f"{end_node.get('name', 'No name')}",
                        'properties': dict(end_node)
                    })
                    node_ids.add(end_node.identity)
                # 处理关系
                edges.append({
                    'from': start_node.identity,
                    'to': end_node.identity,
                    'label': rel.__class__.__name__,
                    'title': f"关系: {rel.__class__.__name__}",
                    'properties': dict(rel)
                })
            return nodes, edges, len(nodes)
        except Exception as e:
            return e